/*
 * @Author: xiangru.xiao xiangru.xiao@mthreads.com
 * @Date: 2023-05-23 20:04:37
 * @LastEditors: xiangru.xiao xiangru.xiao@mthreads.com
 * @LastEditTime: 2023-05-24 18:18:42
 * Copyright (c) 2023 by Moore Threads, All Rights Reserved. 
 */

#include <string>
#include <vector>
#include "objectfactory.h"
#include "node.h"
#include "cu_helper.h"
#include "kernel.cuh"
#define EPS 1e-6
#define TEST_ITERS 4
struct params
{
    int length;
};
bool testResult(cublasHandle_t handle, float* h_a, float* h_b, float* h_c, 
                float* d_a, float* d_b, float* d_c, int length);

class GEMM_test : public Node
{
public:
    GEMM_test()
    {
        nodeUint = "GFLOPS";
    }
    std::string name() override
    {
        return "gemm-test";
    }
    std::string describe() override
    {
        return "矩阵乘法计算";
    }
    Node::NodeType  nodeType() override
    {
        return NodeTP_DataSpeed;
    }
protected:
    bool startup() override;
    bool perform() override;
    void shutdown() override;
private:
    std::vector<params> paramList() const;
private:
    int        deviceid;
};

bool GEMM_test::startup()
{
    deviceid = deviceinfo->deviceID();
    CHECK_CUDA_ERROR(cudaSetDevice(deviceid));
    return true;
}

bool GEMM_test::perform()
{
    float totalValue = 0.0;
    std::vector<params> ls = paramList();
    char* strNum = std::getenv("GEMM_TEST_LEVEL");
    int caseNum = ls.size();
    if(strNum)
    {
        caseNum = std::min(caseNum, atoi(strNum));
    }
    for (int i = 0; i < caseNum; ++i) {
        int length = ls[i].length;
        float alpha = 1.1, beta = 0.9;
        const int real_bytes = sizeof(float) * length * length;
        float *h_a, *h_b, *h_c, *d_a, *d_b, *d_c;
        cublasHandle_t handle;
        CHECK_CUDA_ERROR(cudaMallocHost(&h_a, real_bytes));
        CHECK_CUDA_ERROR(cudaMallocHost(&h_b, real_bytes));
        CHECK_CUDA_ERROR(cudaMallocHost(&h_c, real_bytes));
        CHECK_CUDA_ERROR(cudaMalloc(&d_a, real_bytes));
        CHECK_CUDA_ERROR(cudaMalloc(&d_b, real_bytes));
        CHECK_CUDA_ERROR(cudaMalloc(&d_c, real_bytes));
        CHECK_CUDA_ERROR(cublasCreate(&handle));
        if(testResult(handle, h_a, h_b, h_c, d_a, d_b, d_c, length))
        { //result correct and then do profile
            float totalTime  = 0;
            {
                CUDAEvent evt(&totalTime);
                for (int j = 0; j < TEST_ITERS; ++j) {
                      CHECK_CUDA_ERROR(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, 
                                        length, length, length, 
                                        &alpha, d_a, length, d_b, length, &beta, d_c, length));
                      CHECK_CUDA_ERROR(cudaDeviceSynchronize());
                }
            }
            CHECK_CUDA_ERROR(cublasDestroy(handle));
            CHECK_CUDA_ERROR(cudaFreeHost(h_a));
            CHECK_CUDA_ERROR(cudaFreeHost(h_b));
            CHECK_CUDA_ERROR(cudaFreeHost(h_c));
            CHECK_CUDA_ERROR(cudaFree(d_a));
            CHECK_CUDA_ERROR(cudaFree(d_b));
            CHECK_CUDA_ERROR(cudaFree(d_c));
            totalValue += 2 * float(length) * float(length) * float(length)
                          * TEST_ITERS / totalTime / 1e6f;
        }
        else
        { //get wrong result
            nodeValue = 0.0;
            CHECK_CUDA_ERROR(cublasDestroy(handle));
            CHECK_CUDA_ERROR(cudaFreeHost(h_a));
            CHECK_CUDA_ERROR(cudaFreeHost(h_b));
            CHECK_CUDA_ERROR(cudaFreeHost(h_c));
            CHECK_CUDA_ERROR(cudaFree(d_a));
            CHECK_CUDA_ERROR(cudaFree(d_b));
            CHECK_CUDA_ERROR(cudaFree(d_c));
            return false;
        }
        
    }
    nodeValue = totalValue;
    return true;
}

void GEMM_test::shutdown()
{

}

std::vector<params> GEMM_test::paramList() const
{
    std::vector<params> list;
    params params_[] = {512, 1024, 2048, 4096, 8192};
    const size_t nSizes =  sizeof(params_) / sizeof (params);
    for (int i = 0; i < nSizes; ++i) {
        list.push_back(params_[i]);
    }
    return list;
}

REGISTER_OBJECT(Node, GEMM_test)

bool testResult(cublasHandle_t handle, float* h_a, float* h_b, float* h_c, 
                float* d_a, float* d_b, float* d_c, int length)
{
    int N = length * length;
    float alpha = 1.25, beta = 0.5;
    float* h_out = (float*)malloc(sizeof(float) * N);
    for(int i = 0; i < N; i++)
    {
        h_a[i] = alpha;
        h_b[i] = beta;
        h_c[i] = float(i%16);
    }
    CHECK_CUDA_ERROR(cudaMemcpy(d_a, h_a, sizeof(float) * N, cudaMemcpyHostToDevice));
    CHECK_CUDA_ERROR(cudaMemcpy(d_b, h_b, sizeof(float) * N, cudaMemcpyHostToDevice));
    CHECK_CUDA_ERROR(cudaMemcpy(d_c, h_c, sizeof(float) * N, cudaMemcpyHostToDevice));
    CHECK_CUDA_ERROR(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, 
                                        length, length, length, 
                                        &alpha, d_a, length, d_b, length, &beta, d_c, length));
    CHECK_CUDA_ERROR(cudaDeviceSynchronize());
    CHECK_CUDA_ERROR(cudaMemcpy(h_out, d_c, sizeof(float) * N, cudaMemcpyDeviceToHost));

    double maxv = alpha * beta * length * alpha + beta * 15;
    double nrmse = 0;  // normalized root mean square error
    for (int i = 0; i < N; i++) {
      double dr = alpha * beta * length * alpha + beta * h_c[i] - h_out[i];
      nrmse += (dr * dr);
    }
    nrmse /= (double)(N);
    nrmse = sqrt(nrmse);
    nrmse /= maxv;
    free(h_out);
    return nrmse < EPS;
}